import os
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_community.chat_models import ChatTongyi
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory

# 1、从环境变量读取API密钥和端点URL
AI_QW_KEY = os.getenv("AI_QW_KEY")  # 获取通义千问的API密钥
AI_QW_URL = os.getenv("AI_QW_URL")  # 获取通义千问的API端点（可选）

# 设置环境变量供Langchain使用
os.environ["DASHSCOPE_API_KEY"] = AI_QW_KEY  # 设置通义千问API密钥
os.environ["DASHSCOPE_ENDPOINT"] = AI_QW_URL  # 设置通义千问API端点（如果需要）

# 初始化大语言模型
# 创建ChatTongyi模型实例，用于与通义千问API交互
model = ChatTongyi()

# 2、定义提示模板
# 创建对话模板，包含系统消息和预留的历史消息位置
prompt_template = ChatPromptTemplate.from_messages([
    ("system", "你是一个乐于助人的助手。用{language}尽你所能回答问题。"),  # 系统角色设定，{language}是参数
    MessagesPlaceholder(variable_name="my_msg")  # 预留位置，用于插入对话历史
])

# 3、创建处理链
# 使用管道操作符(|)将提示模板、模型和输出解析器连接起来
# 这是一个LangChain表达式语言(LCEL)的链式调用
chain = prompt_template | model  # 输入先经过提示模板处理，然后传递给模型

# 4、管理聊天历史记录
store = {}  # 用于存储所有用户的聊天历史，键为session_id，值为ChatMessageHistory对象

def get_session_history(session_id: str):
    """根据session_id获取或创建聊天历史记录"""
    if session_id not in store:
        store[session_id] = ChatMessageHistory()  # 为新session创建空的历史记录
    return store[session_id]

# 创建带有消息历史管理功能的可运行对象
do_message = RunnableWithMessageHistory(
    chain,  # 基础处理链
    get_session_history,  # 获取历史记录的调用函数
    input_messages_key='my_msg'  # 指定输入中消息的键名
)

# 配置会话ID，用于区分不同用户的聊天历史
config = {'configurable': {'session_id': 'lhy'}}  # 使用'lhy'作为当前会话的标识符

# 第一轮对话
resp1 = do_message.invoke(
    {
        'my_msg': [HumanMessage(content='你好啊！ 我是李浩洋')],  # 用户消息
        'language': '中文'  # 模板参数，指定使用语言
    },
    config=config  # 传入会话配置
)
print(resp1.content)  # 输出模型回复

# 第二轮对话（模型能记住历史）
resp2 = do_message.invoke(
    {
        'my_msg': [HumanMessage(content='请问：我的名字是什么？')],  # 新用户消息
        'language': '中文'
    },
    config=config  # 使用相同的session_id以保持对话连续性
)
print(resp2.content)  # 输出模型回复（应该能记住名字）

# 第三轮对话：使用流式输出
config = {'configurable': {'session_id': 'lhy'}}  # 同样的会话ID
# 使用stream方法获取流式响应（逐个token输出）
for resp in do_message.stream(
    {
        'my_msg': [HumanMessage(content='我今天生日写一首诗送给我')],
        'language': '中文'
    },
    config=config
):
    # 每次迭代返回一个token，end='-'表示用连字符分隔输出（实际使用时通常不需要）
    print(resp.content, end='-')  # 流式输出每个标记